package websocket

import (
	"fmt"
	"sync"

	"github.com/gorilla/websocket"
)

// 连接池结构体
type WebSocketPool struct {
	clients map[string]*websocket.Conn // 连接存储
	lock    sync.RWMutex               // 读写锁保证并发安全
}

var pool = NewWebSocketPool() // 创建 WebSocket 连接池

// 创建新的 WebSocket 连接池
func NewWebSocketPool() *WebSocketPool {
	return &WebSocketPool{
		clients: make(map[string]*websocket.Conn),
	}
}

// 添加连接到池
func (pool *WebSocketPool) AddClient(clientID string, conn *websocket.Conn) {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	pool.clients[clientID] = conn
	fmt.Println("新增客户端:", clientID)
}

// 从池中移除连接
func (pool *WebSocketPool) RemoveClient(clientID string) {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	if conn, ok := pool.clients[clientID]; ok {
		conn.Close() // 关闭连接
		delete(pool.clients, clientID)
		fmt.Println("移除客户端:", clientID)
	}
}

// 发送消息给指定客户端
func (pool *WebSocketPool) SendMessage(clientID string, message string) error {
	pool.lock.RLock()
	defer pool.lock.RUnlock()

	if conn, ok := pool.clients[clientID]; ok {
		return conn.WriteMessage(websocket.TextMessage, []byte(message))
	}
	return fmt.Errorf("客户端 %s 不存在", clientID)
}

// 广播消息给所有客户端
func (pool *WebSocketPool) BroadcastMessage(message string) {
	pool.lock.RLock()
	defer pool.lock.RUnlock()
	for clientID, conn := range pool.clients {
		err := conn.WriteMessage(websocket.TextMessage, []byte(message))
		if err != nil {
			fmt.Println("发送消息失败:", clientID, err)
		}
	}
}
